# Definition for singly-linked list.
# class ListNode:
#     def __init__(self, x):
#         self.val = x
#         self.next = None

class ListNode:
    def __init__(self, x):
        self.val = x
        self.next = None

class ListBuilder:
    def build(self, numbers):
        head = None
        previous = None
        for i in range(len(numbers)):
            node = ListNode(numbers[i])
            if i == 0:
                head = node
            elif previous:
                previous.next = node
            previous = node
        return head

class Solution:
    def removeNthFromEnd(self, head, n):
        """
        :type head: ListNode
        :type n: int
        :rtype: ListNode
        """
        nthPreviousNode = None
        currentNode = head
        count = 1
        while currentNode.next:
            currentNode = currentNode.next
            count += 1
            if count > n:
                if nthPreviousNode == None:
                    nthPreviousNode = head
                else: 
                    nthPreviousNode = nthPreviousNode.next
        
        if nthPreviousNode:
            nthPreviousNode.next = nthPreviousNode.next.next
        elif count == n:
            head = head.next
        return head

class Test:
    builder = ListBuilder()
    solution = Solution()
    
    def main(self):
        self.testNormal()
        self.testNthFromEndIsFirstOne()
        print("ok")

    def log(self, head):
        desc = "%d" % head.val
        current = head
        while current.next:
            current = current.next
            desc += " -> %d" % (current.val)
        print(desc)
        return desc

    def isEqual(self, head1, head2):
        return self.log(head1) == self.log(head2)

    def testNormal(self):
        head = self.builder.build([1, 2, 3, 4])
        expectedHead = self.builder.build([1, 2, 4])
        result = self.solution.removeNthFromEnd(head, 2)
        assert(self.isEqual(result, expectedHead))

    def testNthFromEndIsFirstOne(self):
        head = self.builder.build([1, 2, 3, 4])
        expectedHead = self.builder.build([2, 3, 4])
        result = self.solution.removeNthFromEnd(head, 4)
        assert(self.isEqual(result, expectedHead))
        
Test().main()




        
